from torch import nn, optim
from torchsummary import summary
from models.cv import ResNet18, ResUnit, ResBlock
from skorch import NeuralNetClassifier


def test_res_unit():
    res_unit = ResUnit(3, 64)
    summary(res_unit, (3, 32, 32))


def test_res_block():
    res_block = ResBlock(3, 64, 2)
    summary(res_block, (3, 32, 32))


def test_resnet18():
    resnet18 = ResNet18()
    summary(resnet18, (3, 32, 32))


def train_resnet18():
    net = NeuralNetClassifier(
        ResNet18,
        max_epochs=10,
        lr=0.001,
        batch_size=128,
        optimizer=optim.Adam,
        optimizer__weight_decay=1e-4,
        iterator_train__shuffle=True,
    )


if __name__ == '__main__':
    test_resnet18()